/*
 * Copyright (C) 2025 ameise <ameise.wang@gmail.com> - All Rights Reserved
 *
 * This file is part of e3ds.
 *
 * e3ds is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * e3ds is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with e3ds. If not, see <https://www.gnu.org/licenses/>.
 */

package heap

import (
	stdHeap "container/heap"
	"gitee.com/ameise84/e3ds/heap/compare"
	"sync"
)

func newHeap[T1 Integer, T2 compare.Ordered[T2]](result compare.Result, size ...int) Heap[T1, T2] {
	n := 32
	if size != nil && size[0] > 0 {
		n = size[0]
	}
	return Heap[T1, T2]{
		base:     newHeapBase[T1, T2](n, result),
		pool:     newPool[T1, T2](),
		entities: map[T1]*entity[T1, T2]{},
	}
}

type Heap[T1 Integer, T2 compare.Ordered[T2]] struct {
	base     heapBase[T1, T2]
	pool     *sync.Pool
	entities map[T1]*entity[T1, T2]
	zeroT1   T1
	zeroT2   T2
}

func (ts *Heap[T1, T2]) PushNoError(ctx T2) T1 {
	tag, _ := ts.Push(ctx)
	return tag
}

func (ts *Heap[T1, T2]) Push(ctx T2) (T1, error) {
	e := ts.pool.Get().(*entity[T1, T2])
	if x, ok := any(ctx).(WithUniqueID[T1]); ok {
		e.id = x.UniqueID()
		if _, ok = ts.entities[e.id]; ok {
			return e.id, ErrHeapPushID
		}
	}
	ts.push(e, ctx)
	return e.id, nil
}

func (ts *Heap[T1, T2]) Update(id T1, ctx T2) bool {
	if e, ok := ts.entities[id]; ok {
		e.ctx = ctx
		stdHeap.Fix(&ts.base, e.index)
		return true
	}
	return false
}

func (ts *Heap[T1, T2]) Peek() (T2, error) {
	if ts.base.Len() == 0 {
		return ts.zeroT2, ErrHeapEmpty
	}
	e := ts.base.entities[0]
	return e.ctx, nil
}

func (ts *Heap[T1, T2]) Pop() (T2, error) {
	if ts.base.Len() == 0 {
		return ts.zeroT2, ErrHeapEmpty
	}
	e := stdHeap.Pop(&ts.base).(*entity[T1, T2])
	delete(ts.entities, e.id)
	ctx := e.ctx
	e.ctx = ts.zeroT2
	ts.pool.Put(e)
	return ctx, nil
}

func (ts *Heap[T1, T2]) Find(id T1) (T2, bool) {
	if v, ok := ts.entities[id]; ok {
		return v.ctx, true
	}
	return ts.zeroT2, false
}

func (ts *Heap[T1, T2]) Remove(id T1) (ctx T2, ok bool) {
	var e *entity[T1, T2]
	if e, ok = ts.entities[id]; ok {
		ctx = e.ctx
		stdHeap.Remove(&ts.base, e.index)
		e.ctx = ts.zeroT2
		delete(ts.entities, id)
		ts.pool.Put(e)
	}
	return
}

func (ts *Heap[T1, T2]) Clean() {
	ts.base.entities = ts.base.entities[:0]
	ts.entities = map[T1]*entity[T1, T2]{}
}

func (ts *Heap[T1, T2]) CleanToSortedSlice() []T2 {
	result := make([]T2, 0, ts.base.Len())
	for {
		item, err := ts.Pop()
		if err != nil {
			break
		}
		result = append(result, item)
	}
	return result
}

func (ts *Heap[T1, T2]) CleanToSlice() []T2 {
	objs := make([]T2, 0, ts.base.Len())
	for _, evt := range ts.base.entities {
		objs = append(objs, evt.ctx)
		evt.ctx = ts.zeroT2
		ts.pool.Put(evt)
	}
	ts.base.entities = ts.base.entities[:0]
	ts.entities = map[T1]*entity[T1, T2]{}
	return objs
}

func (ts *Heap[T1, T2]) CleanToMap() map[T1]T2 {
	objs := make(map[T1]T2, ts.base.Len())
	for _, evt := range ts.base.entities {
		objs[evt.id] = evt.ctx
		evt.ctx = ts.zeroT2
		ts.pool.Put(evt)
	}
	ts.base.entities = ts.base.entities[:0]
	ts.entities = map[T1]*entity[T1, T2]{}
	return objs
}

func (ts *Heap[T1, T2]) Range(f func(T2) bool) {
	for _, v := range ts.base.entities {
		if !f(v.ctx) {
			break
		}
	}
}

func (ts *Heap[T1, T2]) Size() int {
	return ts.base.Len()
}

func (ts *Heap[T1, T2]) IsEmpty() bool {
	return ts.base.Len() == 0
}

func (ts *Heap[T1, T2]) push(evt *entity[T1, T2], ctx T2) {
	evt.ctx = ctx
	stdHeap.Push(&ts.base, evt)
	ts.entities[evt.id] = evt
}

func (ts *Heap[T1, T2]) PushBatch(items []T2) ([]T1, error) {
	result := make([]T1, 0, len(items))
	for _, item := range items {
		id, err := ts.Push(item)
		if err != nil {
			return nil, err
		}
		result = append(result, id)
	}
	return result, nil
}

func (ts *Heap[T1, T2]) PopBatch(n int) ([]T2, error) {
	if n <= 0 || ts.IsEmpty() {
		return nil, ErrHeapEmpty
	}
	if n > ts.Size() {
		return nil, ErrHeapNotEnough
	}

	result := make([]T2, 0, n)
	for i := 0; i < n; i++ {
		if item, err := ts.Pop(); err == nil {
			result = append(result, item)
		}
	}
	return result, nil
}

func (ts *Heap[T1, T2]) RemoveBatch(ids []T1) map[T1]T2 {
	result := make(map[T1]T2)
	for _, id := range ids {
		if item, ok := ts.Remove(id); ok {
			result[id] = item
		}
	}
	return result
}

func (ts *Heap[T1, T2]) Sort(items []T2) ([]T2, error) {
	ts.Clean()

	for _, item := range items {
		if _, err := ts.Push(item); err != nil {
			return nil, ErrHeapPushID
		}
	}

	result := make([]T2, 0, len(items))
	for {
		item, err := ts.Pop()
		if err != nil {
			break
		}
		result = append(result, item)
	}
	return result, nil
}

func (ts *Heap[T1, T2]) ToSortedSlice() []T2 {
	temp := make([]*entity[T1, T2], ts.base.Len())
	copy(temp, ts.base.entities)

	result := ts.CleanToSortedSlice()

	ts.base.entities = temp
	for _, e := range ts.base.entities {
		ts.entities[e.id] = e
	}

	return result
}
